# -*- coding: utf-8 -*-
'''
Created on 23.08.2019

@author: yu03
'''

import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_interpolation_2, FFT_interpolation_compare, FFT_cal
import numpy as np
import os
from scipy.optimize import curve_fit
from Video_Experiment import file_name

def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return np.concatenate(( np.zeros(int((N)/2)), (cumsum[N:]-cumsum[:-N])/float(N) , np.zeros(int((N-1)/2)) ))
def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q

def fit_func_line(x, p, q):
    return p*x+q

fs_cam = 50

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_angle_centers = []
ver_angle_centers = []
hor_phase_centers = []
ver_phase_centers = []
hor_freq_set = []
ver_freq_set = []
hor_m_k_num_set = []
ver_m_k_num_set = []
hor_f_fit_set = []
ver_f_fit_set = []
hor_phi_fit_set = []
ver_phi_fit_set = []

cap = cv2.VideoCapture(file_name)
frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(frame_num, "Frames")
print('ok')
# while(cap.isOpened()):
for k in range(frame_num):
    ret, frame = cap.read() 
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    img_set.append(img)
#     cv2.imshow('frame',img)
#     if cv2.waitKey(1) & 0xFF == ord('q'):
#         break
# cap.release()
# cv2.destroyAllWindows()

# cv2.imshow('image',img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

hor_threshold, ver_threshold = 30, 30
hor_centerline, ver_centerline = img_set[int(frame_num/2)][int(len(img_set[0])/2), :], img_set[int(frame_num/2)][:, int(len(img_set[0][1])/2)]
hor_start, ver_start = np.where(hor_centerline > hor_threshold)[0][0], np.where(ver_centerline > ver_threshold)[0][0]
hor_end, ver_end = np.where(hor_centerline > hor_threshold)[0][-1], np.where(ver_centerline > ver_threshold)[0][-1]
num = 0
for img in img_set[::]:
    num += 1
#     hor_center = img[572][434:834]
#     ver_center = img[:,634][372:772]
    
#     hor_center = img[int((ver_start+ver_end)/2)][hor_start:hor_end]
#     ver_center = img[:,int((hor_start+hor_end)/2)][ver_start:ver_end]
    
    hor_center = np.array(img[int((ver_start+ver_end)/2)])
    ver_center = img[:,int((hor_start+hor_end)/2)]
    
#     hor_center = np.diff(hor_center.astype('float32'))
#     hor_center = np.concatenate(([0], hor_center))
#     ver_center = np.diff(ver_center.astype('float32'))
#     ver_center = np.concatenate(([0], ver_center))
    
    DC_num = 1000
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
    ver_freq_estim, ver_phase_estim, ver_freqline, ver_sig_magnitude, ver_sig_phase, ver_m_k_num, ver_X_m_k, ver_freq_for_phase = FFT_interpolation_2(ver_center, pix_size, 1e5, DC_num)
    hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.2)[0][0]+DC_num
    hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.2)[0][-1]+DC_num
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]

    hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
    hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
    [hor_a, hor_b, hor_c] = hor_params[0]
    hor_f_fit = hor_b
    hor_f_fit_set.append(hor_f_fit)   
    hor_params = curve_fit(fit_func_phase, hor_fit_x, np.unwrap(hor_fit_phase))
    [hor_q] = hor_params[0]
    hor_phi_fit = p*hor_freqline[hor_m_k_num]+hor_q
    hor_phi_fit_set.append(hor_phi_fit+2*np.pi)
    
    
    ver_FFT_start = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][0]+DC_num
    ver_FFT_end = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][-1]+DC_num
    ver_fit_x = ver_freqline[ver_FFT_start:ver_FFT_end+1]
    ver_fit_freq = ver_sig_magnitude[ver_FFT_start:ver_FFT_end+1]
    ver_params = curve_fit(fit_func, ver_fit_x, ver_fit_freq)
    [ver_a, ver_b, ver_c] = ver_params[0]
    ver_f_fit = ver_b
    ver_f_fit_set.append(ver_f_fit)
    
    print(num, hor_m_k_num, ver_m_k_num)
    hor_freq_set.append(hor_freq_estim)
    ver_freq_set.append(ver_freq_estim)
    
    hor_m_k_num_set.append(hor_m_k_num)
    ver_m_k_num_set.append(ver_m_k_num)
    
    hor_phase_centers.append(hor_phase_estim)
    ver_phase_centers.append(ver_phase_estim)

hor_f_fit_set = np.array(hor_f_fit_set)
ver_f_fit_set = np.array(ver_f_fit_set)
hor_phase_centers = np.array(hor_phase_centers)
ver_phase_centers = np.array(ver_phase_centers)

hor_angle = (V_x-Lamda*hor_f_fit_set/2)*1e6 ### urad
hor_length = np.unwrap(hor_phase_centers)/4/np.pi*Lamda*1e9 ### nm
ver_angle = (V_y-Lamda*ver_f_fit_set/2)*1e6 ### urad
ver_length = np.unwrap(ver_phase_centers)/4/np.pi*Lamda*1e9 ### nm
    
if 1:### Plotting Noise_PSD
    zoom_range = [0*fs_cam, 10*fs_cam]
     
    plt.figure('Tilt')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (hor_angle-np.average(hor_angle))[range[0]:range[1]]
    plt.plot(x[int(zoom_range[0]):int(zoom_range[1])], y[int(zoom_range[0]):int(zoom_range[1])], color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (ver_angle-np.average(ver_angle))[range[0]:range[1]]
    plt.plot(x[int(zoom_range[0]):int(zoom_range[1])], y[int(zoom_range[0]):int(zoom_range[1])], color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle, fs_cam, nperseg=len(hor_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_angle, fs_cam, nperseg=len(ver_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    
    
    #### Length
    plt.figure('Length')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (hor_length-np.average(hor_length))[range[0]:range[1]]
    plt.plot(x[int(zoom_range[0]):int(zoom_range[1])], y[int(zoom_range[0]):int(zoom_range[1])], color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = (ver_length-np.average(ver_length))[range[0]:range[1]]
    plt.plot(x[int(zoom_range[0]):int(zoom_range[1])], y[int(zoom_range[0]):int(zoom_range[1])], color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_length, fs_cam, nperseg=len(hor_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 1) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_length, fs_cam, nperseg=len(ver_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 1) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    plt.show()
    
    
if 0:### Plotting Nonlinearity
    ############### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
    start_cutting, stop_cutting = start_cutting/fs_cam, stop_cutting/fs_cam
    zoom_time = start_cutting, stop_cutting
    
    start_cutting, stop_cutting = 1.6, 6.2 ### in second   
#     zoom_time = [1.2, 1.4]
    
    start_cutting, stop_cutting = int(start_cutting*fs_cam), int(stop_cutting*fs_cam)
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
    selected_time = (np.arange(stop_cutting-start_cutting)+start_cutting)/fs_cam    
    zoom_range = [zoom_time[0]*fs_cam, zoom_time[1]*fs_cam]
     
    plt.figure('Tilt')
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_angle = hor_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    hor_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_angle)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_angle = ver_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    ver_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_angle_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    plt.plot(x, y, color='red', marker=' ')
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_angle_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Tilting")
    plt.ylabel("Vertical Angle (urad)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')

#     ax5 = plt.subplot(4,2,5)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = hor_angle_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Horizontal Tilting")
#     plt.ylabel("Horizontal Angle (urad)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')
#        
#     ax6 = plt.subplot(4,2,6)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = ver_angle_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Vertical Tilting")
#     plt.ylabel("Vertical Angle (urad)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')

    ax7 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle, fs_cam, nperseg=len(hor_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax8 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_angle, fs_cam, nperseg=len(ver_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Tilting Noise_PSD")
    plt.ylabel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-4, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    
    
    #### Length
    plt.figure('Length')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = plt.subplot(3,2,1)
    range = [0, len(hor_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_length = hor_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    hor_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    range = [0, len(ver_length)]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length
    plt.plot(x, y, color='blue', marker=' ')
    
    measured_length = ver_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    ver_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_time, linear_fit, color='red')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = hor_length_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    plt.plot(x, y, color='red', marker=' ')
    range = [start_cutting, stop_cutting]
    x = (np.arange(range[1]-range[0])+range[0])/fs_cam
    y = ver_length_nonlinearity
    plt.plot(x, y, color='blue', marker=' ')
    plt.title("Vertical Length")
    plt.ylabel("Vertical Length (nm)")
    plt.xlabel("Time (s)")
    plt.grid(which='both', axis='both')
    
#     ax5 = plt.subplot(4,2,5)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = hor_length_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Horizontal Length")
#     plt.ylabel("Horizontal Length (nm)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')
#        
#     ax6 = plt.subplot(4,2,6)
#     range = [int(zoom_range[0]-start_cutting), int(zoom_range[1]-start_cutting)]
#     x = (np.arange(zoom_range[1]-zoom_range[0])+zoom_range[0])/fs_cam
#     y = ver_length_nonlinearity[range[0]:range[1]]
#     plt.plot(x, y, color='blue', marker=' ')
#     plt.title("Vertical Length")
#     plt.ylabel("Vertical Length (nm)")
#     plt.xlabel("Time (s)")
#     plt.grid(which='both', axis='both')

    ax7 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_length, fs_cam, nperseg=len(hor_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-3, 2e3) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax8 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(ver_length, fs_cam, nperseg=len(ver_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Vertical Length Noise_PSD")
    plt.ylabel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-3, 2e3) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    plt.tight_layout()
    plt.show()